{
 "cells": [
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "import time\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import random\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from sklearn.model_selection import KFold\n",
    "from sklearn import metrics\n",
    "from utils import load_data, build_graph, weight_reset\n",
    "from model import MAHN\n",
    "import dgl"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-05-08T01:02:50.773448Z",
     "start_time": "2024-05-08T01:02:47.823962400Z"
    }
   },
   "id": "d1bc4f171071931e",
   "execution_count": 1
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "def Train(device, directory, epochs, attn_size, attn_heads, out_dim, dropout, slope, lr, wd, random_seed, sample_num, model_type):\n",
    "    print('sample_num',sample_num)\n",
    "    print('dropout',dropout)\n",
    "    random.seed(random_seed)\n",
    "    np.random.seed(random_seed)\n",
    "    torch.manual_seed(random_seed)\n",
    "    g_D0, g_ME0, g_D1, g_ME1, g, g0, disease_vertices, metabolite_vertices, ID, IME, IMI, samples, \\\n",
    "    DMI_associations, MIME_associations = build_graph(directory, random_seed,device, sample_num)\n",
    "    samples_df = pd.DataFrame(samples, columns=['disease', 'metabolite', 'label'])\n",
    "\n",
    "    print('## vertices:', g.number_of_nodes())\n",
    "    print('## edges:', g.number_of_edges())\n",
    "    print('## disease nodes:', torch.sum(g.ndata['type'] == 1).cpu().numpy())\n",
    "    print('## metabolite nodes: ', torch.sum(g.ndata['type'] == 0).cpu().numpy())\n",
    "    print('## microbe nodes: ', torch.sum(g0.ndata['type'] == 2).cpu().numpy())\n",
    "\n",
    "    auc_result = []\n",
    "    acc_result = []\n",
    "    pre_result = []\n",
    "    recall_result = []\n",
    "    f1_result = []\n",
    "    prc_result = []\n",
    "    fprs = []\n",
    "    tprs = []\n",
    "    precisions = []\n",
    "    recalls = []\n",
    "\n",
    "    # 设置五折交叉验证\n",
    "    i = 0\n",
    "    kf = KFold(n_splits=5, shuffle=True, random_state=random_seed)\n",
    "    for train_idx, test_idx in kf.split(samples[:, 2]):\n",
    "        i += 1\n",
    "        print('Training for Fold', i)\n",
    "        # 将训练集的标记为1，(测试集)其余为0\n",
    "        samples_df['train'] = 0\n",
    "        samples_df['train'].iloc[train_idx] = 1\n",
    "        train_tensor = torch.from_numpy(samples_df['train'].values.astype('int64')).to(device)\n",
    "        edge_data = {'train': train_tensor}\n",
    "        # 对g和g0中疾病和metabolite的两个异质图的训练边进行标记\n",
    "        g.edges[disease_vertices, metabolite_vertices].data.update(edge_data)\n",
    "        g.edges[metabolite_vertices, disease_vertices].data.update(edge_data)\n",
    "\n",
    "        g0.edges[disease_vertices, metabolite_vertices].data.update(edge_data)\n",
    "        g0.edges[metabolite_vertices, disease_vertices].data.update(edge_data)\n",
    "\n",
    "        train_eid = g.filter_edges(lambda edges: edges.data['train']).to(device)\n",
    "        g_train = g.edge_subgraph(train_eid, preserve_nodes=True)\n",
    "        g_train = g_train.to(device)\n",
    "        g_train0 = g0.edge_subgraph(train_eid, preserve_nodes=True)\n",
    "        g_train0 = g_train0.to(device)\n",
    "\n",
    "        # 训练标签和测试标签用于最后对比\n",
    "        label_train = g_train.edata['label'].unsqueeze(1).to(device)\n",
    "        src_train, dst_train = g_train.all_edges()\n",
    "\n",
    "        test_eid = g.filter_edges(lambda edges: edges.data['train'] == 0).to(device)\n",
    "        src_test, dst_test = g.find_edges(test_eid)\n",
    "        label_test = g.edges[test_eid].data['label'].unsqueeze(1).to(device)\n",
    "        print('## Training edges:', len(train_eid))\n",
    "        print('## Testing edges:', len(test_eid))\n",
    "\n",
    "        if model_type == 'MAHN':\n",
    "            model = MAHN(G_D0=g_D0, G_D1=g_D1, G_ME0=g_ME0, G_ME1=g_ME1,\n",
    "                           hidden_dim=64, G=g_train0,\n",
    "                           meta_paths_list=['dme', 'med', 'dmi', 'mime'],\n",
    "                           feature_attn_size=attn_size,\n",
    "                           num_heads=attn_heads,\n",
    "                           num_diseases=ID.shape[0],\n",
    "                           num_metabolite=IME.shape[0],\n",
    "                           num_microbe=IMI.shape[0],\n",
    "                           d_sim_dim=ID.shape[1],\n",
    "                           me_sim_dim=IME.shape[1],\n",
    "                           mi_sim_dim=IMI.shape[1],\n",
    "                           out_dim=out_dim,\n",
    "                           dropout=dropout,\n",
    "                           slope=slope, device=device\n",
    "                           )\n",
    "            model.apply(weight_reset)\n",
    "            model.to(device)\n",
    "            optimizer = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=wd)\n",
    "            loss = nn.BCELoss()  # Binary Cross Entropy Loss二元交叉熵损失函数\n",
    "\n",
    "            best_val_auc = 0\n",
    "            best_epoch = 0\n",
    "            for epoch in range(epochs):\n",
    "                start = time.time()\n",
    "                model.train()  # 模型训练\n",
    "\n",
    "                with torch.autograd.set_detect_anomaly(True):\n",
    "                    g_D0 = g_D0.to(device)\n",
    "                    g_D1 = g_D1.to(device)\n",
    "                    g_ME0 = g_ME0.to(device)\n",
    "                    g_ME1 = g_ME1.to(device)\n",
    "                    g_train0 = g_train0.to(device)\n",
    "                    g0 = g0.to(device)\n",
    "                    src_train = src_train.to(device)\n",
    "                    dst_train = dst_train.to(device)\n",
    "\n",
    "                    score_train = model(g_D0, g_D1, g_ME0, g_ME1, g_train0, g0, src_train, dst_train).to(device)\n",
    "                    loss_train = loss(score_train, label_train).to(device)\n",
    "\n",
    "                    optimizer.zero_grad()\n",
    "                    loss_train.backward()\n",
    "                    optimizer.step()\n",
    "\n",
    "                model.eval()  # 模型验证\n",
    "                with torch.no_grad():\n",
    "                    score_val = model(g_D0, g_D1, g_ME0, g_ME1, g, g0, src_test, dst_test).to(device)\n",
    "                    loss_val = loss(score_val, label_test).to(device)\n",
    "\n",
    "                score_train_cpu = np.squeeze(score_train.detach().cpu().numpy())\n",
    "                score_val_cpu = np.squeeze(score_val.detach().cpu().numpy())\n",
    "                label_train_cpu = np.squeeze(label_train.detach().cpu().numpy())\n",
    "                label_val_cpu = np.squeeze(label_test.detach().cpu().numpy())\n",
    "\n",
    "                train_auc = metrics.roc_auc_score(label_train_cpu, score_train_cpu)\n",
    "                val_auc = metrics.roc_auc_score(label_val_cpu, score_val_cpu)\n",
    "\n",
    "                pred_val = [0 if j < 0.5 else 1 for j in score_val_cpu]\n",
    "                acc_val = metrics.accuracy_score(label_val_cpu, pred_val)\n",
    "                pre_val = metrics.precision_score(label_val_cpu, pred_val)\n",
    "                recall_val = metrics.recall_score(label_val_cpu, pred_val)\n",
    "                f1_val = metrics.f1_score(label_val_cpu, pred_val)\n",
    "                end = time.time()\n",
    "                if (epoch + 1) % 10 == 0:\n",
    "                    print('Epoch:', epoch + 1, 'Train Loss: %.4f' % loss_train.item(),\n",
    "                      'Val Loss: %.4f' % loss_val.cpu().detach().numpy(),\n",
    "                      'Acc: %.4f' % acc_val, 'Pre: %.4f' % pre_val, 'Recall: %.4f' % recall_val, 'F1: %.4f' % f1_val,\n",
    "                      'Train AUC: %.4f' % train_auc, 'Val AUC: %.4f' % val_auc, #'Time: %.2f' % (end - start)\n",
    "                      )\n",
    "\n",
    "                if val_auc > best_val_auc:\n",
    "                    best_val_auc = val_auc\n",
    "                    best_epoch = epoch + 1\n",
    "                    torch.save(model.state_dict(), 'best_model.pth')  # 保存每一折的最优模型\n",
    "\n",
    "            model.load_state_dict(torch.load('best_model.pth'))  # 重载最优模型进行测试\n",
    "            model.eval()  # 模型测试\n",
    "            with torch.no_grad():\n",
    "                score_test = model(g_D0, g_D1, g_ME0, g_ME1, g, g0, src_test, dst_test).to(device)\n",
    "\n",
    "            score_test_cpu = np.squeeze(score_test.detach().cpu().numpy())\n",
    "\n",
    "            label_test_cpu = np.squeeze(label_test.detach().cpu().numpy())\n",
    "\n",
    "            fpr, tpr, thresholds = metrics.roc_curve(label_test_cpu, score_test_cpu)\n",
    "            precision, recall, _ = metrics.precision_recall_curve(label_test_cpu, score_test_cpu)\n",
    "            test_auc = metrics.auc(fpr, tpr)\n",
    "            test_prc = metrics.auc(recall, precision)\n",
    "\n",
    "            pred_test = [0 if j < 0.5 else 1 for j in score_test_cpu]\n",
    "            acc_test = metrics.accuracy_score(label_test_cpu, pred_test)\n",
    "            pre_test = metrics.precision_score(label_test_cpu, pred_test)\n",
    "            recall_test = metrics.recall_score(label_test_cpu, pred_test)\n",
    "            f1_test = metrics.f1_score(label_test_cpu, pred_test)\n",
    "\n",
    "            print('Fold:', i, 'Best Epoch:', best_epoch, 'Test acc: %.4f' % acc_test, 'Test Pre: %.4f' % pre_test,\n",
    "                  'Test Recall: %.4f' % recall_test, 'Test F1: %.4f' % f1_test, 'Test PRC: %.4f' % test_prc,\n",
    "                  'Test AUC: %.4f' % test_auc)\n",
    "\n",
    "            auc_result.append(test_auc)\n",
    "            acc_result.append(acc_test)\n",
    "            pre_result.append(pre_test)\n",
    "            recall_result.append(recall_test)\n",
    "            f1_result.append(f1_test)\n",
    "            prc_result.append(test_prc)\n",
    "\n",
    "            fprs.append(fpr)\n",
    "            tprs.append(tpr)\n",
    "            precisions.append(precision)\n",
    "            recalls.append(recall)\n",
    "\n",
    "    print('## Training Finished !')\n",
    "    print('-----------------------------------------------------------------------------------------------')\n",
    "    print('Auc', auc_result)\n",
    "    print('Acc', acc_result)\n",
    "    print('Pre', pre_result)\n",
    "    print('Recall', recall_result)\n",
    "    print('F1', f1_result)\n",
    "    print('Prc', prc_result)\n",
    "    print('AUC mean: %.4f, variance: %.4f \\n' % (np.mean(auc_result), np.std(auc_result)),\n",
    "          'Accuracy mean: %.4f, variance: %.4f \\n' % (np.mean(acc_result), np.std(acc_result)),\n",
    "          'Precision mean: %.4f, variance: %.4f \\n' % (np.mean(pre_result), np.std(pre_result)),\n",
    "          'Recall mean: %.4f, variance: %.4f \\n' % (np.mean(recall_result), np.std(recall_result)),\n",
    "          'F1-score mean: %.4f, variance: %.4f \\n' % (np.mean(f1_result), np.std(f1_result)),\n",
    "          'PRC mean: %.4f, variance: %.4f \\n' % (np.mean(prc_result), np.std(prc_result)))\n",
    "    print('fprs', fprs)\n",
    "    print('tprs', tprs)\n",
    "    print('precisions', precisions)\n",
    "    print('recalls', recalls)\n",
    "    return fprs, tprs, auc_result, precisions, recalls, prc_result"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-05-08T01:02:50.810448800Z",
     "start_time": "2024-05-08T01:02:50.796447800Z"
    }
   },
   "id": "9ca39d3abb57695f",
   "execution_count": 2
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
